package baidu

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"quipus/internal/llm"

	"github.com/baidubce/bce-qianfan-sdk/go/qianfan"
	mark3labs "github.com/mark3labs/mcp-go/mcp"
)

type QianfanClient struct {
	client *qianfan.ChatCompletion
	model  string
}

type qianfanAuthInfo struct {
	SecretKey string `json:"secret_key"`
	AccessKey string `json:"access_key"`
}

func NewClient(config *llm.LLMClientConfig) (*QianfanClient, error) {
	authInfo := qianfanAuthInfo{}
	err := json.Unmarshal([]byte(config.AuthInfo), &authInfo)
	if err != nil {
		return nil, err
	}

	qianfan.GetConfig().AccessKey = authInfo.AccessKey
	qianfan.GetConfig().SecretKey = authInfo.SecretKey
	client := qianfan.NewChatCompletion()
	c := &QianfanClient{
		client: client,
		model:  config.Model,
	}
	return c, nil
}

func (c *QianfanClient) Chat(ctx context.Context, req *llm.ChatReq) (err error) {
	if req == nil {
		return nil
	}

	model := req.Model
	msgs := req.Msgs

	if len(msgs) < 1 {
		return errors.New("message is null")
	}

	if len(model) < 1 {
		model = c.model
	}

	if len(model) < 1 {
		model = "ERNIE-Speed-8K"
	}

	messages := []qianfan.ChatCompletionMessage{}
	systemPrompt := ""
	for _, m := range msgs {
		if m.Role == "system" {
			systemPrompt = m.Content
			continue
		}

		messages = append(messages, qianfan.ChatCompletionMessage{
			Role:    m.Role,
			Content: m.Content,
		})
	}
	c.client.Model = model

	stream, err := c.client.Stream(context.TODO(),
		&qianfan.ChatCompletionRequest{Messages: messages, System: systemPrompt})

	if err != nil {
		panic(err)
	}
	defer stream.Close()
loop:
	for {
		r, err := stream.Recv()
		if err != nil {
			panic(err)
		}
		if r.IsEnd {
			break
		}
		select {
		case req.RspChan <- r.Result:
			// slog.Info("get response: ", r.Result)
		case <-ctx.Done():
			slog.Info("baidu recv quit signal")
			break loop
		}

	}
	req.CompleteChan <- true
	return
}

func (c *QianfanClient) GetProviderName() string {
	return "qianfan"
}

func (c *QianfanClient) Complete(ctx context.Context, req *llm.CompleteReq) (llm.Message, error) {
	if req == nil {
		return nil, nil
	}

	model := req.Model
	msgs := req.Msgs

	if len(msgs) < 1 {
		return nil, errors.New("message is null")
	}

	if len(model) < 1 {
		model = "ERNIE-Speed-8K"
	}

	messages := []qianfan.ChatCompletionMessage{}
	systemPrompt := ""
	for _, m := range msgs {
		if m.Role == "system" {
			systemPrompt = m.Content
			continue
		}
		messages = append(messages, qianfan.ChatCompletionMessage{
			Role:    m.Role,
			Content: m.Content,
		})
	}
	c.client.Model = model
	chatReq := &qianfan.ChatCompletionRequest{Messages: messages, System: systemPrompt}
	if req.Tools != nil && len(req.Tools) > 0 {
		chatReq.Functions = convertTool(req.Tools)
	}

	res, err := c.client.Do(ctx, chatReq)
	if err != nil {
		return nil, err
	}
	return &WeiXinMessage{Resp: res, Role: llm.RoleAssistant}, nil
}

func convertTool(tools []mark3labs.Tool) []qianfan.Function {
	var zpTools []qianfan.Function
	for _, tool := range tools {
		t := qianfan.Function{
			Name:        tool.Name,
			Description: tool.Description,
			Parameters:  convertSchema(tool.InputSchema),
		}
		zpTools = append(zpTools, t)
	}

	return zpTools
}
func convertSchema(schema mark3labs.ToolInputSchema) map[string]interface{} {
	required := schema.Required
	if required == nil {
		required = []string{}
	}

	return map[string]interface{}{
		"type":       schema.Type,
		"properties": schema.Properties,
		"required":   required,
	}
}

type WeiXinMessage struct {
	Resp *qianfan.ModelResponse
	Role string
}

func (m *WeiXinMessage) GetRole() string {
	return llm.RoleAssistant
}

func (m *WeiXinMessage) GetContent() string {
	return m.Resp.Result
}

func (m *WeiXinMessage) GetToolCalls() []llm.ToolCall {
	var calls []llm.ToolCall
	calls = append(calls, &ToolCallWrapper{m.Resp.FunctionCall})
	return calls
}

func (m *WeiXinMessage) IsToolResponse() bool {
	return m.Role == llm.RoleTool
}

func (m *WeiXinMessage) GetToolResponseID() string {
	return m.Resp.FunctionCall.Name
}

func (m *WeiXinMessage) GetUsage() (int, int) {
	return int(m.Resp.Usage.CompletionTokens), int(m.Resp.Usage.PromptTokens)
}

type ToolCallWrapper struct {
	Call *qianfan.FunctionCall
}

func (t *ToolCallWrapper) GetID() string {
	return t.Call.Name
}

func (t *ToolCallWrapper) GetName() string {
	return t.Call.Name
}

func (t *ToolCallWrapper) GetArguments() map[string]interface{} {
	var args map[string]interface{}
	if err := json.Unmarshal([]byte(t.Call.Arguments), &args); err != nil {
		return make(map[string]interface{})
	}
	return args
}
